#include <chrono>
#include <iostream>
#include <thread>

#include <asio_driver/device/buzzer_src.hpp>
#include <asio_driver/device/rc_sbus.hpp>
#include <asio_driver/utils/range.hpp>
#include <asio_driver_special/imu_wit_cmp10a.hpp>
#include <asio_driver_special/motor_brt_can.hpp>
#include <asio_driver_special/motor_dc_pwm.hpp>
#include <ctrl_alg/filter/window_filter.hpp>
#include <ctrl_alg/pid/pid_pos.hpp>

#include <general_interface/msg/chassis_control.hpp>
#include <general_interface/msg/chassis_state.hpp>
#include <rclcpp/rclcpp.hpp>

static const int gyro_rate = 200;
static const int angle_rate = 50;
static const int speed_rate = 10;
static const int state_publish_rate = 50;

/*
 * params:
 * gyro_pid: [p, i, d], default [200, 0, 0], PID parameter for gyro control
 * angle_pid: [p, i, d], default [5, 0, 0], PID parameter for angle control
 * speed_pid: [p, i, d], default [10, 0.1, 0], PID parameter for speed control
 * rotate_pid: [p, i, d], default [40, 0, 0], PID parameter for rotate control
 * */

class Devices
{
  private:
    bool isOpened = false;

  public:
    AD::AsioContext::SharePtr context;
    AD::Special::MotorBrtCan motorLeft;
    AD::Special::MotorBrtCan motorRight;
    AD::Special::MotorBrtCan motorExternalLeft;
    AD::Special::MotorBrtCan motorExternalRight;
    AD::Special::IMUWitCmp10a imu;
    AD::Buzzer buzzer;

  public:
    Devices(AD::AsioContext::SharePtr context)
        : context(context), motorLeft(context), motorRight(context), motorExternalLeft(context),
          motorExternalRight(context), imu(context), buzzer(context)
    {
    }
    ~Devices()
    {
        if (isOpened)
        {
            Close();
        }
    }

    void Open()
    {
        std::unique_ptr<AD::Can> can = std::make_unique<AD::Can>(context);
        AD::CanBus::SharePtr can_bus = AD::CanBus::Create(context);
        motorLeft.Open(can_bus, AD::Special::MotorBrtCan::MotorConfig{.address = 4, .poleCount = 14});
        motorRight.Open(can_bus, AD::Special::MotorBrtCan::MotorConfig{.address = 3, .poleCount = 14});
        motorExternalLeft.Open(can_bus, AD::Special::MotorBrtCan::MotorConfig{.address = 1, .poleCount = 14});
        motorExternalRight.Open(can_bus, AD::Special::MotorBrtCan::MotorConfig{.address = 2, .poleCount = 14});
        can->Open("can0");
        can_bus->Open(std::move(can));
        imu.Open("/dev/ttyS6", 115200);
        buzzer.Open(34, true);
        isOpened = true;
    }

    void Close()
    {
        motorLeft.Close();
        motorRight.Close();
        motorExternalLeft.Close();
        motorExternalRight.Close();
        imu.Close();
        buzzer.Close();
        isOpened = false;
    }
};

struct State
{
    double gyrox;
    double gyroy;
    double gyroz;
    double accx;
    double accy;
    double accz;
    double angley;
    double anglex;
    double anglez;
    double speed;

    double forwardCurrent;
    double rotateCurrent;
    double leftCurrent;
    double rightCurrent;

    double leftRPM;
    double rightRPM;

    double targetGyro;
    double targetAngle;
    double targetSpeed;
    double targetRotate;
};

class ChassisNode : public rclcpp::Node
{
  private:
    enum class ControlMode
    {
        Current,
        Gyro,
        Angle,
        Speed
    };

  private:
    AD::AsioContext::SharePtr context;
    std::unique_ptr<Devices> devices;
    State state;
    std::unique_ptr<CtrlAlg::PIDPos<double>> gyroPID;
    std::unique_ptr<CtrlAlg::PIDPos<double>> anglePID;
    std::unique_ptr<CtrlAlg::PIDPos<double>> speedPID;
    std::unique_ptr<CtrlAlg::PIDPos<double>> rotatePID;
    rclcpp::Subscription<general_interface::msg::ChassisControl>::SharedPtr sub_chassis_control;
    rclcpp::Publisher<general_interface::msg::ChassisState>::SharedPtr pub_chassis_state;
    rclcpp::TimerBase::SharedPtr timer_balance;
    rclcpp::TimerBase::SharedPtr timer_angle;
    rclcpp::TimerBase::SharedPtr timer_speed;
    rclcpp::TimerBase::SharedPtr timer_beep;
    rclcpp::TimerBase::SharedPtr timer_publish_state;
    CtrlAlg::WindowFilter<double> speedFilter = CtrlAlg::WindowFilter<double>(10);
    ControlMode mode = ControlMode::Speed;

  public:
    ChassisNode() : Node("chassis_ros2")
    {
        RCLCPP_INFO(get_logger(), "init chassis devices");
        context = AD::AsioContext::CreatAsioContext();
        devices = std::make_unique<Devices>(context);
        devices->Open();
        RCLCPP_INFO(get_logger(), "chassis devices inited");

        RCLCPP_INFO(get_logger(), "load parameters for PID");
        gyroPID = CreatPIDWithParam<CtrlAlg::PIDPos<double>>("gyro", 200, 0, 0);
        anglePID = CreatPIDWithParam<CtrlAlg::PIDPos<double>>("angle", 5, 0, 0);
        speedPID = CreatPIDWithParam<CtrlAlg::PIDPos<double>>("speed", 10, 0.1, 0);
        speedPID->SetMaxIntegralOutputPart(3);
        rotatePID = CreatPIDWithParam<CtrlAlg::PIDPos<double>>("rotate", 40, 0, 0);
        RegeistParamCallBack();
        RCLCPP_INFO(get_logger(), "PID parameters loaded");

        RCLCPP_INFO(get_logger(), "regeist publisher and subscriber");
        sub_chassis_control = create_subscription<general_interface::msg::ChassisControl>(
            "chassis_control", 1, std::bind(&ChassisNode::OnChassisControl, this, std::placeholders::_1));
        pub_chassis_state = create_publisher<general_interface::msg::ChassisState>("chassis_state", 1);
        RCLCPP_INFO(get_logger(), "chassis_control subscriber registered");

        RCLCPP_INFO(get_logger(), "creat timer for control loop");
        timer_balance =
            create_wall_timer(std::chrono::milliseconds(1000 / gyro_rate), std::bind(&ChassisNode::OnGyroLoop, this));
        timer_angle =
            create_wall_timer(std::chrono::milliseconds(1000 / angle_rate), std::bind(&ChassisNode::OnAngleLoop, this));
        timer_speed =
            create_wall_timer(std::chrono::milliseconds(1000 / speed_rate), std::bind(&ChassisNode::OnSpeedLoop, this));
        timer_beep = create_wall_timer(std::chrono::milliseconds(500), std::bind(&ChassisNode::OnBeepFinished, this));
        timer_publish_state = create_wall_timer(std::chrono::milliseconds(1000 / state_publish_rate),
                                                std::bind(&ChassisNode::OnPublishState, this));
        RCLCPP_INFO(get_logger(), "timer for control loop created");
        Beep();
    }

    void Beep()
    {
        if (!timer_beep->is_canceled())
        {
            timer_beep->execute_callback();
            timer_beep->cancel();
        }
        devices->buzzer.Beep();
        timer_beep->reset();
    }

  private:
    template <class T> std::unique_ptr<T> CreatPIDWithParam(std::string name, double p, double i, double d)
    {
        std::vector<double> values = {p, i, d};
        rcl_interfaces::msg::ParameterDescriptor descriptor;
        descriptor.description = "PID parameter for " + name;
        declare_parameter<std::vector<double>>(name + "_pid", values, descriptor);
        values = get_parameter(name + "_pid").as_double_array();
        return std::make_unique<CtrlAlg::PIDPos<double>>(values[0], values[1], values[2]);
    }

    void RegeistParamCallBack()
    {
        auto _ = add_on_set_parameters_callback(
            [&](const std::vector<rclcpp::Parameter> &parameters)
            {
                rcl_interfaces::msg::SetParametersResult result;
                result.successful = true;
                for (const auto &param : parameters)
                {
                    if (param.get_name() == "gyro_pid")
                    {
                        auto values = param.as_double_array();
                        gyroPID->SetPID(values[1], values[2], values[3]);
                    }
                    else if (param.get_name() == "angle_pid")
                    {
                        auto values = param.as_double_array();
                        anglePID->SetPID(values[1], values[2], values[3]);
                    }
                    else if (param.get_name() == "speed_pid")
                    {
                        auto values = param.as_double_array();
                        speedPID->SetPID(values[1], values[2], values[3]);
                    }
                    else if (param.get_name() == "rotate_pid")
                    {
                        auto values = param.as_double_array();
                        rotatePID->SetPID(values[1], values[2], values[3]);
                    }
                    else
                    {
                        result.successful = false;
                        result.reason = "unknown parameter name: " + param.get_name();
                        break;
                    }
                }
                return result;
            });
    }

    void OnGyroLoop()
    {
        context->SpinOnce();

        auto imuData = devices->imu.Read();

        imuData.gyro[1] = -imuData.gyro[1];
        imuData.angle[1] = -(imuData.angle[1] + 1);

        state.gyrox = AD::Utils::InRange(imuData.gyro[0], -30, 30) ? imuData.gyro[0] : state.gyrox;
        state.gyroy = AD::Utils::InRange(imuData.gyro[1], -30, 30) ? imuData.gyro[1] : state.gyroy;
        state.gyroz = AD::Utils::InRange(imuData.gyro[2], -30, 30) ? imuData.gyro[2] : state.gyroz;
        state.accx = AD::Utils::InRange(imuData.acc[0], -100000, 100000) ? imuData.acc[0] : state.accx;
        state.accy = AD::Utils::InRange(imuData.acc[1], -100000, 100000) ? imuData.acc[1] : state.accy;
        state.accz = AD::Utils::InRange(imuData.acc[2], -100000, 100000) ? imuData.acc[2] : state.accz;
        state.anglex = AD::Utils::InRange(imuData.angle[0], -30, 30) ? imuData.angle[0] : state.anglex;
        state.angley = AD::Utils::InRange(imuData.angle[1], -30, 30) ? imuData.angle[1] : state.angley;
        state.anglez = AD::Utils::InRange(imuData.angle[2], -30, 30) ? imuData.angle[2] : state.anglez;

        if (mode != ControlMode::Current)
        {
            state.forwardCurrent = gyroPID->GetControl(state.gyroy - state.targetGyro, 5);
            state.forwardCurrent = AD::Utils::ClampAbs(state.forwardCurrent, 5000);
        }
        state.rotateCurrent = rotatePID->GetControl(state.gyroz - state.targetRotate, 5);
        state.rotateCurrent = AD::Utils::ClampAbs(state.rotateCurrent, 1000);

        state.leftCurrent = AD::Utils::ClampAbs(state.forwardCurrent - state.rotateCurrent, 6000);
        state.rightCurrent = AD::Utils::ClampAbs(state.forwardCurrent + state.rotateCurrent, 6000);

        devices->motorLeft.SendHeartBeat();
        devices->motorRight.SendHeartBeat();
        devices->motorLeft.SetCurrent(state.leftCurrent);
        devices->motorRight.SetCurrent(state.rightCurrent);
    }

    void OnAngleLoop()
    {
        if (mode != ControlMode::Gyro)
        {
            state.targetGyro = -anglePID->GetControl(state.angley - state.targetAngle, 20);
        }
    }

    void OnSpeedLoop()
    {
        devices->motorLeft.RequestRPM();
        devices->motorRight.RequestRPM();

        state.leftRPM = devices->motorLeft.ReadState().rpm;
        state.rightRPM = devices->motorRight.ReadState().rpm;
        state.speed = (state.leftRPM + state.rightRPM) / 2;
        state.speed = speedFilter.GetControl(state.speed);
        if (mode != ControlMode::Angle)
        {
            state.targetAngle = speedPID->GetControl(state.targetSpeed - state.speed, 100);
            state.targetAngle = AD::Utils::ClampAbs(state.targetAngle, 10);
        }
    }

    void OnPublishState()
    {
        auto msg = general_interface::msg::ChassisState();
        auto vec3 = geometry_msgs::msg::Vector3();
        vec3.x = state.gyrox;
        vec3.y = state.gyroy;
        vec3.z = state.gyroz;
        msg.angular_velocity = vec3;
        vec3.x = (state.leftRPM + state.rightRPM) / 2;
        msg.linear_velocity = vec3;
        vec3.x = state.anglex;
        vec3.y = state.angley;
        vec3.z = state.anglez;
        msg.angle = vec3;
        pub_chassis_state->publish(msg);
    }

    void OnBeepFinished() { devices->buzzer.StopBeep(); }

    void OnChassisControl(const general_interface::msg::ChassisControl::SharedPtr msg)
    {
        state.targetGyro = msg->angular;
        switch (msg->mode)
        {
        case general_interface::msg::ChassisControl::CURRENT:
            mode = ControlMode::Current;
            state.forwardCurrent = msg->velocity;
            break;
        case general_interface::msg::ChassisControl::GYRO:
            mode = ControlMode::Gyro;
            state.targetGyro = msg->velocity;
            break;
        case general_interface::msg::ChassisControl::ANGLE:
            mode = ControlMode::Angle;
            state.targetAngle = msg->velocity;
            break;
        case general_interface::msg::ChassisControl::SPEED:
            mode = ControlMode::Speed;
            state.targetSpeed = msg->velocity;
            break;
        default:
            RCLCPP_WARN(get_logger(), "unknown control mode: %d", msg->mode);
            break;
        }
    }
};

int main(int argc, char **argv)
{
    rclcpp::init(argc, argv);
    auto chassis_node = std::make_shared<ChassisNode>();
    rclcpp::spin(chassis_node);
    rclcpp::shutdown();
    return 0;
}